fix: raise ValueError on xp=np + jnp-backed-grid mismatch - #337
Merged
Conversation
…Maker Adds a guard at the top of AbstractMaker.__init__ that raises ValueError when the caller declared xp=np but handed in a jnp-backed grid (grid.use_jax=True). This is almost always a user error inside an @jax.jit body — they forgot to pass xp=jnp to the library method. Without this guard, the function body silently host-transfers the tracer through __array__() (slow) or fails at the JIT boundary because the if xp is np: guard wraps the result in an aa.Array2D. The new explicit error points users at the lens_calc.py workspace guide for the JIT-it-yourself pattern. Final PR of Phase 2 of z_features/jax_user_intro.md (autoarray/simulator_use_jax.md). Design doc: admin_jammy/notes/jax_interface.md §4.8 Issue: PyAutoArray#334 Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Final PR of Phase 2 of
z_features/jax_user_intro.md. Adds a clear, loudValueErroratAbstractMaker.__init__when the caller passesxp=npbut the input grid is jnp-backed (grid.use_jax=True). This is almost always a user error: someone wrote@jax.jit def f(grid): ...; return some_method(grid)and forgot to passxp=jnptosome_method.Today's silent failure modes:
__array__()on the JAX tracer → host transfer (slow, "works" but secretly wrong).if xp is np:guard inside library functions wraps the result inaa.Array2D, which then fails at the JIT boundary with a confusingTypeError: ... is not a valid JAX type.With this PR, the guard fires at the function-entry boundary with a clear message pointing at the
lens_calc.pyworkspace guide (Phase 5d, where the JIT-it-yourself pattern is canonically documented).API Changes
AbstractMaker.__init__now raisesValueErrorwhenxp is npandgrid.use_jaxisTrue. Pre-existing call sites are unaffected — all of PyAutoArray, PyAutoLens, PyAutoGalaxy continue to pass the tests (verified: 837 / 317 / 918 pass).See full details below.
Test Plan
test_autoarray/structures/decorators/test_abstract_xp_mismatch.py:xp=np+ jnp-backed grid (with the expected error message format).xp=np+ NumPy grid (preserves existing behaviour).use_jaxattribute (e.g. plain ndarray).Full API Changes
Changed behaviour
AbstractMaker.__init__(func, obj, grid, xp=np, *args, **kwargs)raisesValueErrorat the top of the function whenxp is npandgetattr(grid, "use_jax", False)isTrue. Error message includes the called function's qualname and a pointer to thelens_calc.pyworkspace guide.Migration
xpand grid types.xp=jnpto the library method call inside your@jax.jitbody, matching the jnp-backed grid you're passing in.Future work
_strict_xp=False) for library internals that legitimately want the host-transfer NumPy fallback path. Not added in this PR because no current site needs it; will add if a real use case surfaces.🤖 Generated with Claude Code